# 引入必要的库
import cv2                          # OpenCV库，用于实现图像处理
import numpy as np                  # NumPy库，用于进行矩阵和数组操作
import matplotlib.pyplot as plt     # Matplotlib的pyplot模块，用于图像的显示和绘制
import pickle

# 加载保存的模型
with open('E:/IPA_projects/stanstone/clf_model.pkl', 'rb') as f:
    clf = pickle.load(f)

# 读取砂岩截面图2及其对应分区图  
original_image = cv2.imread('E:/IPA_projects/stanstone/Sandstone_imgs/Sandstone_2.tif', 0)  
segment = cv2.imread('E:/IPA_projects/stanstone/Sandstone_imgs/Sandstone_2_segment.tif', 0)  
  
# 获取图像数据和标签  
X = original_image.reshape(-1,1)  

# 对新图像进行预测并显示结果  
y_pred = clf.predict(X)
segmentation = y_pred.reshape(original_image.shape[:2])

# 计算准确率
accuracy = np.mean(segmentation == segment)

# 使用matplotlib显示原分割图和转换后的图像  
fig, ax = plt.subplots(1, 3, figsize=(16, 6))    
ax[0].imshow(original_image,cmap='gray')    
ax[0].set_title("original image", fontsize=15)    
ax[1].imshow(segment,cmap='gray')    
ax[1].set_title("segment", fontsize=15)    
ax[2].imshow(segmentation,cmap='gray')    
ax[2].set_title("segmentation(acc:{:.3f})".format(accuracy), fontsize=15)    
plt.tight_layout()    
plt.show()